import asyncio
import json
import time

import websockets
from abc import ABC, abstractmethod
from websockets.legacy.server import WebSocketServerProtocol


class WebSocketClient:
    def __init__(self, uri: str):
        """
        初始化 WebSocket 客户端
        :param uri: WebSocket 服务器地址（例如 "ws://localhost:8765"）
        """
        self.uri = uri
        self.websocket = None

    async def connect(self):
        """
        连接到 WebSocket 服务器
        """
        if not self.websocket:
            try:
                self.websocket = await websockets.connect(self.uri, max_size=1024 * 1024 * 4)
                print("已连接到 WebSocket 服务器")
            except Exception as e:
                exit("websocket连接失败！")

    async def send(self, message):
        """
        发送消息到 WebSocket 服务器
        :param message: 要发送的消息
        """
        if not self.websocket:
            raise ConnectionError("未连接到服务器，请先调用 connect 方法")
        await self.websocket.send(message)
        print(f"已发送消息")

    async def response_msg_handler(self):
        async for response in self.websocket:
            print(response)
            if "status" in response:
                break

    async def close(self):
        """
        关闭 WebSocket 连接
        """
        if self.websocket:
            await self.websocket.close()
            self.websocket = None
            print("WebSocket 连接已关闭")


class WebSocketServer(ABC):
    def __init__(self):
        try:
            asyncio.run(self.start_server(self.receive_msg_handler))
        except KeyboardInterrupt:
            print("服务器已关闭")

    @abstractmethod
    # 消息处理函数
    async def receive_msg_handler(self, websocket: WebSocketServerProtocol):
        async for message in websocket:
            print(f"收到消息: {message}")
            # 处理消息

    # 启动 WebSocket 服务器
    @staticmethod
    async def start_server(receive_msg_handler=None):
        async with websockets.serve(
                receive_msg_handler,
                "0.0.0.0",  # 允许远程访问
                8765,  # WebSocket 端口
                ping_interval=60,  # 保持连接活跃
                max_size=1024 * 1024 * 4
        ):
            print("WebSocket Server Started at 0.0.0.0:8765")
            await asyncio.Future()


if __name__ == "__main__":
    client = WebSocketClient(uri="ws://124.88.174.116:31919/")

    async def main():
        # 连接到服务器
        await client.connect()
        # 连续发送两条消息，无延时
        s = time.time()
        await client.send(json.dumps({"input": "你好，你知道我是谁吗，你说这个是什么意思"}, ensure_ascii=False))
        await client.response_msg_handler()
        print(time.time() - s)
        # await client.send('{"name":"zs"}')
        # await client.response_msg_handler()

        # 关闭连接
        await client.close()

    asyncio.run(main())